package org.taj.adapter.llama;

import com.alibaba.fastjson2.JSON;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.taj.adapter.server.Session;
import org.zeromq.ZMQ;

import java.util.List;

@Component
public class LLaMAClient {
    private OkHttpClient client = new OkHttpClient();
    @Value("${llama.host}")
    private String host;
    @Value("${llama.port}")
    private String port;

    public void send(Session session, ZMQ.Socket pubSocket, String sessionId) {
        StringBuilder respBuffer = new StringBuilder();
        RequestBody body = RequestBody.create(session.toString(), MediaType.parse("application/json"));
        Request req = new Request.Builder().url("http://" + host + ":" + port + "/v1/chat/completions").post(body).build();
        EventSource.Factory factory = EventSources.createFactory(client);
        factory.newEventSource(req, new EventSourceListener() {
            @Override
            public void onOpen(EventSource eventSource, Response response) {
                System.out.println("SSE连接已开启");
            }

            @Override
            public void onEvent(EventSource eventSource, String id, String type, String data) {
                if (data.equals("[DONE]")) {
                    eventSource.cancel();
                } else {
                    LLaMAResponse response = JSON.parseObject(data, LLaMAResponse.class);
                    List<LLaMAResponse.Choice> choices = response.getChoices();
                    LLaMAResponse.Choice choice = choices.get(0);
                    if (choice.getFinishReason() != null && choice.getFinishReason().equals("stop")) {
                        session.addAssistantResponse(respBuffer.toString());
                        if(response.getTimings() != null) {
                            LLaMAResponse.Timing timings = response.getTimings();
                            long promptN = timings.getPromptN();
                            double promptMs = timings.getPromptMs();
                            long predictedN = timings.getPredictedN();
                            double predictedMs = timings.getPredictedMs();
                            session.setTimings(promptN, promptMs, predictedN, predictedMs);
                        }
                        pubSocket.send(sessionId);
                        System.out.println();
                    } else {
                        String content = choice.getDelta().getContent();
                        respBuffer.append(content);
                        pubSocket.sendMore(sessionId);
                        pubSocket.send(content);
                        System.out.print(content);
                    }
                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                System.out.println("SSE连接已关闭");
            }

            @Override
            public void onFailure(EventSource eventSource, Throwable t, Response response) {
                System.err.println("SSE连接失败: " + t.getMessage());
                if (response != null) {
                    System.err.println("Response Code: " + response.code());
                }
            }
        });

    }
}
